from rknnlite.api import RKNNLite
import cv2
import numpy as np
import multiprocessing as mp

def load_rknn_model(model_path, device_id):
    rknn = RKNNLite()
    ret = rknn.load_rknn(model_path)
    if ret != 0:
        raise RuntimeError(f"Failed to load RKNN model: {ret}")
    
    # 打印模型信息（调试用）
    # print(f"Model {model_path} inputs:", rknn.inputs)
    # print(f"Model {model_path} outputs:", rknn.outputs)
    
    ret = rknn.init_runtime(core_mask=RKNNLite.NPU_CORE_0 << device_id)
    if ret != 0:
        raise RuntimeError(f"Failed to init RKNN runtime: {ret}")
    return rknn

def run_model(model, img_path, task):
    img = cv2.imread(img_path)
    if img is None:
        raise FileNotFoundError(f"Image not found: {img_path}")
    
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    img = cv2.resize(img, (640, 640))  # 确保与模型输入尺寸一致
    
    # 转换为 4D 输入 (1, C, H, W)
    img = np.transpose(img, (2, 0, 1))  # HWC -> CHW
    img = np.expand_dims(img, 0)       # CHW -> NCHW (N=1)
    
    # 尝试两种归一化方式（选择一种）
    img = img.astype(np.float32) / 255.0      # 方式1: [0, 1]
    # img = (img - 127.5) / 128.0             # 方式2: [-1, 1]
    
    # 推理
    outputs = model.inference(inputs=[img])
    print(f"{task} Results Shape:", [arr.shape for arr in outputs])  # 打印输出形状
    print(f"{task} Results Sample:", outputs[0][0, 0, :5])  # 打印前5个值（调试用）

if __name__ == "__main__":
    pose_model = load_rknn_model("./models/merge_pose_epoch_600.rknn", 0)
    seg_model = load_rknn_model("./models/merge_seg_epoch_201.rknn", 1)

    # 多进程并行推理
    p1 = mp.Process(target=run_model, args=(pose_model, "./data/Color/1.png", "Pose"))
    p2 = mp.Process(target=run_model, args=(seg_model, "./data/Color/1.png", "Seg"))
    
    p1.start()
    p2.start()
    p1.join()
    p2.join()